import json
import math
from pathlib import Path
from typing import Dict

from common.chinese_char_dataset import *

ROOT = (Path(__file__) / '../../').resolve()
DATASET = (ROOT / 'dataset').resolve()

DEFAULT_TRANS_PROB = 2e-8


def predict(*, raw_input: str, chn_data: ChineseCharDataset, markov_matrix: Dict[str, Dict[str, float]]) -> str:
    raw_input = raw_input.strip()

    if raw_input == '':
        return ''

    input = raw_input.strip().split(' ')

    for pinyin in input:
        if not chn_data.is_pinyin_in_table(pinyin):
            raise ValueError(f'用户所输入的拼音 {pinyin} 不在表中。')

    initial_matrix, transition_matrix, end_matrix = markov_matrix['initial_matrix'], markov_matrix['transition_matrix'], \
                                                    markov_matrix['end_matrix']

    last_prob, curr_prob, last_result, curr_result = {}, {}, {}, {}
    for chn in chn_data.get_chns_of_pinyin(input[0]):
        last_prob[chn] = math.log10(initial_matrix.get(chn, DEFAULT_TRANS_PROB))
        last_result[chn] = chn

    # Viterbi 算法
    for i in range(1, len(input)):
        for curr_chn in chn_data.get_chns_of_pinyin(input[i]):
            for last_chn, last_chn_prob in last_prob.items():
                # 获取两个汉字之间的转移概率；如果找不到，就用 DEFAULT_TRANS_PROB 代替
                trans_prob = transition_matrix.get(f'{last_chn}{curr_chn}', DEFAULT_TRANS_PROB)
                accumulated_prob = last_chn_prob + math.log10(trans_prob)

                if curr_chn not in curr_prob or accumulated_prob > curr_prob[curr_chn]:
                    curr_prob[curr_chn] = accumulated_prob
                    curr_result[curr_chn] = last_result[last_chn] + curr_chn

        last_prob, last_result = curr_prob, curr_result
        curr_prob, curr_result = {}, {}

    for final_chn in list(last_prob.keys()):
        final_prob = end_matrix.get(final_chn, DEFAULT_TRANS_PROB)
        last_prob[final_chn] += math.log10(final_prob)

    sorted_probs = [last_result[x] for x in sorted(last_prob.keys(), key=lambda k: last_prob[k], reverse=True)]
    return sorted_probs[0]


def load_dataset():
    print('正在加载中文拼音数据集...')
    chn_data = ChineseCharDataset(char_table=(DATASET / 'chinese-char-table.txt').read_text(encoding='utf-8-sig'))

    print('正在加载马尔可夫转移矩阵...')
    with (DATASET / 'markov-matrix.json').open('r', encoding='utf-8-sig') as f:
        markov_matrix = json.load(f)

    print('加载完毕。\n')

    return chn_data, markov_matrix

def handle_input_txt(chn_data, markov_matrix):
    raw_inputs = (DATASET / 'input.txt').read_text(encoding='utf-8-sig').strip().split('\n')
    result = []

    for raw_input in raw_inputs:
        print(f'正在将 {raw_input} 转换为汉字...')

        result.append(predict(
            raw_input=raw_input,
            chn_data=chn_data,
            markov_matrix=markov_matrix,
        ))

    (DATASET / 'output.txt').write_text('\n'.join(result), encoding='utf-8-sig')

def test() -> None:
    chn_data, markov_matrix = load_dataset()

    print('直接按下回车，即可读取 input.txt 文件。')
    while True:
        raw_input = input('> ')
        if raw_input == '':
            print('\n未输入任何内容，正在读取 input.txt 文件中的输入……\n')
            handle_input_txt(chn_data, markov_matrix)
            break

        try:
            print(predict(
                raw_input=raw_input,
                chn_data=chn_data,
                markov_matrix=markov_matrix,
            ), end='\n\n')
        except Exception as e:
            print('错误：' + str(e), end='\n\n')



if __name__ == '__main__':
    test()

